#include <stdio.h>
#include <stdlib.h>
#include <numa.h>
#include <numaif.h>

int main() {
    // Initialize NUMA
    if (numa_available() < 0) {
        printf("NUMA not available\n");
        return 1;
    }

    // Get number of NUMA nodes
    int num_nodes = numa_num_configured_nodes();
    printf("Number of NUMA nodes: %d\n", num_nodes);

    // Get the node mask to check available nodes
    struct bitmask *node_mask = numa_get_mems_allowed();
    printf("Available NUMA nodes: ");
    for (int i = 0; i < num_nodes; i++) {
        if (numa_bitmask_isbitset(node_mask, i)) {
            printf("%d ", i);
        }
    }
    printf("\n");

    // Allocate memory on different NUMA nodes
    size_t size = 1024 * 1024; // 1MB
    void *ptr1, *ptr2, *ptr3, *ptr4;
    
    // Allocate on specific nodes with error checking
    ptr1 = numa_alloc_onnode(size, 0);
    if (ptr1 == NULL) {
        perror("Failed to allocate on node 0");
        return 1;
    }

    ptr2 = NULL;
    if (num_nodes > 1) {
        ptr2 = numa_alloc_onnode(size, 1);
        if (ptr2 == NULL) {
            perror("Failed to allocate on node 1");
            numa_free(ptr1, size);
            return 1;
        }
    }

    ptr3 = NULL;
    if (num_nodes > 2) {
        ptr3 = numa_alloc_onnode(size, 2);
        if (ptr3 == NULL) {
            perror("Failed to allocate on node 2");
            numa_free(ptr1, size);
            if (ptr2) numa_free(ptr2, size);
            return 1;
        }
    }

    ptr4 = NULL;
    if (num_nodes > 3) {
        ptr4 = numa_alloc_onnode(size, 3);
        if (ptr4 == NULL) {
            perror("Failed to allocate on node 3");
            numa_free(ptr1, size);
            if (ptr2) numa_free(ptr2, size);
            if (ptr3) numa_free(ptr3, size);
            return 1;
        }
    }

    // Add memory touching to ensure pages are actually allocated
    char *p;
    if (ptr1) {
        for (p = ptr1; p < (char*)ptr1 + size; p += 4096) *p = 0;
    }
    if (ptr2) {
        for (p = ptr2; p < (char*)ptr2 + size; p += 4096) *p = 0;
    }
    if (ptr3) {
        for (p = ptr3; p < (char*)ptr3 + size; p += 4096) *p = 0;
    }
    if (ptr4) {
        for (p = ptr4; p < (char*)ptr4 + size; p += 4096) *p = 0;
    }

    // Check which NUMA node the memory belongs to
    int node;
    if (ptr1) {
        get_mempolicy(&node, NULL, 0, ptr1, MPOL_F_NODE | MPOL_F_ADDR);
        printf("Memory at %p belongs to NUMA node: %d\n", ptr1, node);
    }

    if (ptr2) {
        get_mempolicy(&node, NULL, 0, ptr2, MPOL_F_NODE | MPOL_F_ADDR);
        printf("Memory at %p belongs to NUMA node: %d\n", ptr2, node);
    }

    if (ptr3) {
        get_mempolicy(&node, NULL, 0, ptr3, MPOL_F_NODE | MPOL_F_ADDR);
        printf("Memory at %p belongs to NUMA node: %d\n", ptr3, node);
    }

    if (ptr4) {
        get_mempolicy(&node, NULL, 0, ptr4, MPOL_F_NODE | MPOL_F_ADDR);
        printf("Memory at %p belongs to NUMA node: %d\n", ptr4, node);
    }

    // Cleanup
    if (ptr1) numa_free(ptr1, size);
    if (ptr2) numa_free(ptr2, size);
    if (ptr3) numa_free(ptr3, size);
    if (ptr4) numa_free(ptr4, size);

    return 0;
}
